################################################################################
# The MIT License (MIT)
#
# Copyright (c) 2014 Kevin Haroldsen
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
################################################################################

import re
import enum
import hivex
import libhivexmod
import struct


class BootMode(enum.Enum):
    BIOS = enum.auto()
    EFI = enum.auto()


class RegValueType(enum.IntEnum):
    REG_NONE = 0
    REG_SZ = 1
    REG_EXPAND_SZ = 2
    REG_BINARY = 3
    REG_DWORD = 4
    REG_DWORD_BIG_ENDIAN = 5
    REG_LINK = 6
    REG_MULTI_SZ = 7
    REG_RESOURCE_LIST = 8
    REG_FULL_RESOURCE_DESCRIPTIOR = 9
    REG_RESOURCE_REQUIREMENTS_LISTS = 10
    REG_QWORD = 11


class HiveNavigator:

    '''# The perfect companion class to hivex, it's a HiveNavigator that allows simple parsing of a registry hive.'''

    def __init__(self, hive, root, nav=None):
        self._hive = hive
        self._root = root
        if nav is None:
            # we are at "/"
            self._nav = [self._root]
        else:
            self._nav = nav[:]

    def __str__(self):
        return '%s(%s)' % (self.__class__.__name__, self.tell())

    def __repr__(self):
        return str(self)

    def exists(self, path):
        '''Returns True if the given pathation (key) exists.'''

        new = self.__class__(self._hive, self._root, nav=self._nav)
        return new._seek(path, False)

    def get(self, path):
        '''Get a new HiveNavigator class focused at path.'''

        new = self.__class__(self._hive, self._root, nav=self._nav)
        new._seek(path, True)
        return new

    def tell(self):
        '''Gives a path representation of the focus in the hive.'''

        return '/' + '/'.join([self._hive.node_name(i) for i in self._nav][1:])

    def seek(self, path):
        '''Seek to a new subkey. Inodes and paths are supported.'''

        self._seek(path, True)

    def key(self):
        '''Returns the name of the current key.'''

        return self._hive.node_name(self._cur())

    def subkeys(self, path='.'):
        '''Lists all subkeys of the current focused key or path.'''

        return [self._hive.node_name(i) for i in self._hive.node_children(self.get(path)._cur())]

    def value(self, name, **kwargs):
        '''Gets the value in a key-value pair in path. default can be used.'''

        nav = self.get(kwargs['path']) if 'path' in kwargs else self
        v = None

        try:
            vn = nav._hive.node_get_value(nav._cur(), name)
            t = nav._hive.value_type(vn)[0]

            if t == RegValueType.REG_SZ:
                v = nav._hive.value_string(vn)
            elif t == RegValueType.REG_MULTI_SZ:
                v = nav._hive.value_multiple_strings(vn)
            elif t == RegValueType.REG_DWORD:
                v = nav._hive.value_dword(vn)
            elif t == RegValueType.REG_QWORD:
                v = nav._hive.value_qword(vn)
            else:
                v = nav._hive.value_value(vn)[1]
                if t == RegValueType.REG_BINARY and isinstance(v, str):
                    v = v.encode()

            if v is None:
                raise RuntimeError

        except RuntimeError:
            if 'default' in kwargs:
                return kwargs['default']
            else:
                raise KeyError('Value "%s" not found in key "%s"' % (name, nav.tell()))

        return v

    def values(self, path='.'):
        '''Lists all values of this key or path.'''

        return [self._hive.value_key(i) for i in self._hive.node_values(self.get(path)._cur())]

    def value_dict(self, path='.'):
        '''Provides a dictionary of the key-value pairs in this registry key.'''

        return {k: self.value(k, path=path) for k in self.values(path)}

    def value_exists(self, name, path='.'):
        '''Returns True if the value name exists in key path.'''

        nav = self.get(path) if path != '.' else self

        try:
            nav._hive.node_get_value(nav._cur(), name)
            return True
        except RuntimeError:
            return False

    def value_type(self, name, path='.'):
        '''Returns the type of the value in path'''

        nav = self.get(path) if path != '.' else self

        try:
            vn = nav._hive.node_get_value(nav._cur(), name)
            return int(nav._hive.value_type(vn)[0])
        except RuntimeError:
            raise KeyError('Value "%s" not found in key "%s"' % (name, nav.tell()))

    def walk(self, path='.'):
        '''Recurses through path and returns (focus, keys, value_dict) similar to os.walk'''

        keys = self.subkeys(path)
        vals = self.value_dict(path)
        yield (path, keys, vals)
        for key in keys:
            newpath = path
            if not (path.endswith('/') or path.endswith('\\')):
                newpath += '/'
            newpath += key
            for i in self.walk(newpath):
                yield i

    def create_key(self, path):
        '''Creates a new key at the specified path. Path can be absolute or relative.'''

        parent_nav = None
        key_name = None
        if True:
            idx = path.rfind('/')
            if idx != -1:
                parent_nav = self.get(path[:idx])
                key_name = path[idx+1:]
            else:
                idx = path.rfind('\\')
                if idx != -1:
                    parent_nav = self.get(path[:idx])
                    key_name = path[idx+1:]
                else:
                    parent_nav = self
                    key_name = path

        new_node = self._hive.node_add_child(parent_nav._cur(), key_name)

        # Create and return a new navigator for the new key
        return self.__class__(self._hive, self._root, nav=parent_nav._nav + [new_node])

    def delete_key(self, path):
        '''Deletes the key at the specified path. Path can be absolute or relative.'''

        if not self.exists(path):
            return

        nav = self.get(path)
        parent_nav = nav.get("..")
        self._hive.node_delete(parent_nav._cur(), nav._cur())

    def set_value(self, name, value_type, value, path='.'):
        '''Sets the value in a key-value pair in path. Create value if neccessary.'''

        nav = self.get(path) if path != '.' else self

        # Set the new value based on type
        if value_type == RegValueType.REG_SZ:
            self._set_value(nav._cur(), name, RegValueType.REG_SZ, value.encode('utf-16-le') + b'\x00\x00')
        elif value_type == RegValueType.REG_EXPAND_SZ:
            assert False
        elif value_type == RegValueType.REG_BINARY:
            self._set_value(nav._cur(), name, RegValueType.REG_BINARY, value)
        elif value_type == RegValueType.REG_DWORD:
            self._set_value(nav._cur(), name, RegValueType.REG_DWORD, struct.pack('<I', value))
        elif value_type == RegValueType.REG_MULTI_SZ:
            self._set_value(nav._cur(), name, RegValueType.REG_MULTI_SZ, b''.join(x.encode('utf-16-le') + b'\x00\x00' for x in value) + b'\x00\x00')
        elif value_type == RegValueType.REG_QWORD:
            assert False
        else:
            assert False

    def _cur(self):
        '''The currently focused inode'''

        return self._nav[-1]

    def _seek(self, path, raise_exception):
        if isinstance(path, int):
            newnav = []
            try:
                while path != self._root:
                    newnav.append(path)
                    path = self._hive.node_parent(path)
                newnav.append(self._root)
                newnav.reverse()
            except RuntimeError:
                if raise_exception:
                    raise ValueError('Could not locate node %d' % path)
                else:
                    return False

            self._nav = newnav
            return True

        if isinstance(path, str):
            newnav = []
            keypath = re.split(r'\\|/', path.rstrip('/\\'))
            if keypath:
                newnav = self._nav[:]
                for key in keypath:
                    if not key:
                        newnav = [self._root]
                    elif key == '.':
                        continue
                    elif key == '..':
                        if len(newnav) > 1:
                            newnav.pop()
                    else:
                        child = self._hive.node_get_child(newnav[-1], key)
                        if not child:
                            if raise_exception:
                                raise ValueError('Could not locate key "%s" in path "%s"' % (key, path))
                            else:
                                return False
                        newnav.append(child)

            self._nav = newnav
            return True

        assert False

    def _set_value(self, node, key, value_type, value):
        self._hive.node_set_value(node, {
            'key': key,
            't': value_type,
            'value': value
        })


def enum(*sequential, **named):
    enums = dict(zip(sequential, range(len(sequential))), **named)
    reverse = {value: key for key, value in enums.items()}
    enumsl = {key.lower(): value for key, value in enums.items()}

    def cast(self, x):
        if isinstance(x, str):
            v = enumsl.get(x.lower())
            if v:
                return v
            else:
                raise KeyError(x + ' is not a valid enum value')
        else:
            return x

    enums['reverse_mapping'] = reverse
    enums['__new__'] = cast
    return type('Enum', (), enums)


def is_guid(val):
    val = str(val)
    if (val.count('{') ^ val.count('}')) or val.count('{') > 1:
        return False
    exp = r'^\{?[0-9A-Fa-f]{8}-([0-9A-Fa-f]{4}-){3}[0-9A-Fa-f]{12}\}?$'
    return bool(re.match(exp, val))


def guid_bracket(val, strict=False):
    if (not strict) or is_guid(val):
        if val.startswith('{'):
            return val
        else:
            return '{' + val + '}'
    else:
        raise ValueError('"%s" is not a valid guid' % val)


def guid_unbracket(val, strict=True):
    if (not strict) or is_guid(val):
        if val.startswith('{'):
            return val[1:-1]
        else:
            return val
    else:
        raise ValueError('"%s" is not a valid guid' % val)


def load_hive(fname, readonly):
    # because hivex craps out with a bad file, but still tries to release in the
    # __del__ method, we first check here to see if it can all be loaded.
    try:
        h = libhivexmod.open(fname, 0)
        libhivexmod.close(h)
    except RuntimeError:
        raise IOError('Unable to load the registry hive "%s"!' % (fname))
    return hivex.Hivex(fname, write=(not readonly))


class callbacklist(list):
    'A list that calls a function when changed'

    callback = None

    def _call(self):
        if callable(self.callback):
            self.callback()

    def __init__(self, *args, **kwargs):
        list.__init__(self, *args, **kwargs)
        for i in ('append', 'clear', 'extend', 'insert', 'pop', 'remove', 'reverse', 'sort'):
            def repl(func):
                def newfunc(*args, **kwargs):
                    v = func(*args, **kwargs)
                    self._call()
                    return v
                return newfunc
            setattr(self, i, repl(getattr(self, i)))


def identity(x):
    return x
